
package decoderstage

import chisel3.util._
import chisel3._
import decoderstage.NPCOpCode._
object NPCOpCode {
  val NPC_Width = 10.U.getWidth
  val NPC_PC4 = 0.U(NPC_Width.W)
  val NPC_J = 1.U(NPC_Width.W)
  val NPC_JAL = 2.U(NPC_Width.W)
  val NPC_JR = 3.U(NPC_Width.W)
  val NPC_JALR = 4.U(NPC_Width.W)
  val NPC_BEQ = 5.U(NPC_Width.W)
  val NPC_BNE = 6.U(NPC_Width.W)
  val NPC_BLEZ = 7.U(NPC_Width.W)
  val NPC_BGTZ = 8.U(NPC_Width.W)
  val NPC_BLTZ = 9.U(NPC_Width.W)
  val NPC_BGEZ = 10.U(NPC_Width.W)
}

class NPC extends Module {
  val io = IO(new Bundle {
    val PC = Input(UInt(32.W))
    val rsData = Input(SInt(32.W))
    val rtData = Input(SInt(32.W))
    val imm16 = Input(UInt(16.W))
    val imm26 = Input(UInt(26.W))
    val NPC_OpCode = Input(UInt(NPC_Width.W))
    val newPC = Output(UInt(32.W))
  })

  val PC = io.PC.asSInt()
  val PC4 = PC + 4.S
  val J_Ans = Cat(PC(31, 28), io.imm26, "b00".U(2.W)).asSInt()
  val JR_Ans = io.rsData
  val BJ_Ans = PC + Cat(io.imm16, "b00".U(2.W)).asSInt()
  val J_Type = io.NPC_OpCode === NPC_J || io.NPC_OpCode === NPC_JAL
  val JR_Type = io.NPC_OpCode === NPC_JR || io.NPC_OpCode === NPC_JALR
  val BJ_OK = MuxCase(false.B, Seq(
    (io.NPC_OpCode === NPC_BEQ) -> (io.rsData === io.rtData),
    (io.NPC_OpCode === NPC_BNE) -> (io.rsData =/= io.rtData),
    (io.NPC_OpCode === NPC_BLEZ) -> (io.rsData <= 0.S),
    (io.NPC_OpCode === NPC_BGTZ) -> (io.rsData > 0.S),
    (io.NPC_OpCode === NPC_BLTZ) -> (io.rsData < 0.S),
    (io.NPC_OpCode === NPC_BGEZ) -> (io.rsData >= 0.S)
  ))
  io.newPC := MuxCase(PC4, Seq(
    J_Type -> J_Ans,
    JR_Type -> JR_Ans,
    BJ_OK -> BJ_Ans
  )).asUInt()
}
